#!/usr/bin/env python
# encoding: utf-8
'''
@author: wangjianrong
@software: pycharm
@file: synthesize_img.py
@time: 2020/10/10 9:45
@desc:
'''

import cv2
import os
import numpy as np
import random
from copy import deepcopy
from opencv_op.read_img import cv_imread

# img = np.ones((1000,1000,3),dtype=np.uint8)*23
# cv2.imshow('img',img)
# cv2.waitKey()

img = cv2.imread('1.tif',-1)
img3 = img[...,:3]
inds = img[...,-1] == 0
cv2.imshow('img3',img)
img[inds] = (20,20,20,20)

cv2.imshow('img4',img)
# cv2.imshow('img3',img3)
# img3[img[...,-1] != 0] = (255,255,255)
# cv2.imshow('img5',img3)
# cv2.imshow('255-img3',255-img3)
cv2.waitKey()
a = img[...,-1]
inds = a == 0
img[inds] = (255,255,255)
cv2.imshow('img',img)
cv2.waitKey()

# print(np.unique(img,return_counts=True))
# print(np.max(img))

class Chunk:
    def __init__(self, file, align=True, bigendian=True, inclheader=False):
        import struct
        self.closed = False
        self.align = align      # whether to align to word (2-byte) boundaries
        if bigendian:
            strflag = '>'
        else:
            strflag = '<'
        self.file = file
        self.chunkname = file.read(4)
        if len(self.chunkname) < 4:
            raise EOFError
        try:
            self.chunksize = struct.unpack_from(strflag+'L', file.read(4))[0]
        except struct.error:
            raise EOFError
        if inclheader:
            self.chunksize = self.chunksize - 8 # subtract header
        self.size_read = 0
        try:
            self.offset = self.file.tell()
        except (AttributeError, OSError):
            self.seekable = False
        else:
            self.seekable = True

    def getname(self):
        """Return the name (ID) of the current chunk."""
        return self.chunkname

    def getsize(self):
        """Return the size of the current chunk."""
        return self.chunksize

    def close(self):
        if not self.closed:
            try:
                self.skip()
            finally:
                self.closed = True

    def isatty(self):
        if self.closed:
            raise ValueError("I/O operation on closed file")
        return False

    def seek(self, pos, whence=0):
        """Seek to specified position into the chunk.
        Default position is 0 (start of chunk).
        If the file is not seekable, this will result in an error.
        """

        if self.closed:
            raise ValueError("I/O operation on closed file")
        if not self.seekable:
            raise OSError("cannot seek")
        if whence == 1:
            pos = pos + self.size_read
        elif whence == 2:
            pos = pos + self.chunksize
        if pos < 0 or pos > self.chunksize:
            raise RuntimeError
        self.file.seek(self.offset + pos, 0)
        self.size_read = pos

    def tell(self):
        if self.closed:
            raise ValueError("I/O operation on closed file")
        return self.size_read

    def read(self, size=-1):
        """Read at most size bytes from the chunk.
        If size is omitted or negative, read until the end
        of the chunk.
        """

        if self.closed:
            raise ValueError("I/O operation on closed file")
        if self.size_read >= self.chunksize:
            return b''
        if size < 0:
            size = self.chunksize - self.size_read
        if size > self.chunksize - self.size_read:
            size = self.chunksize - self.size_read
        data = self.file.read(size)
        self.size_read = self.size_read + len(data)
        if self.size_read == self.chunksize and \
           self.align and \
           (self.chunksize & 1):
            dummy = self.file.read(1)
            self.size_read = self.size_read + len(dummy)
        return data

    def skip(self):
        """Skip the rest of the chunk.
        If you are not interested in the contents of the chunk,
        this method should be called so that the file points to
        the start of the next chunk.
        """

        if self.closed:
            raise ValueError("I/O operation on closed file")
        if self.seekable:
            try:
                n = self.chunksize - self.size_read
                # maybe fix alignment
                if self.align and (self.chunksize & 1):
                    n = n + 1
                self.file.seek(n, 1)
                self.size_read = self.size_read + n
                return
            except OSError:
                pass
        while self.size_read < self.chunksize:
            n = min(8192, self.chunksize - self.size_read)
            dummy = self.read(n)
            if not dummy:
                raise EOFError

import base64

# with open('dh4--92_2020-8-27-9-49-31_00088.iff','rb') as f:
#     # base64_encode = base64.b64encode(f.read()).decode('utf-8')
#
#     data = Chunk(f)
#     print(data.getname())
#     print(data.getsize())
#     base64_encode = data.read()
#     base64_decode = base64.b64decode(base64_encode)
#     img_array = np.fromstring(base64_decode, np.uint8)
#     img = cv2.imdecode(img_array, cv2.COLOR_BGR2RGB)
#     cv2.imshow('img', img)
#     cv2.waitKey()
#
#
# cap = cv2.VideoCapture('明亮.gif')


def filter_smoke_data(img,thresh=15):
    '''
    过滤烟雾图片中黑色块与过小的烟雾
    :return:
    '''
    alpha = img[...,-1]
    h,w = alpha.shape
    contours,_ = cv2.findContours(alpha.copy(),cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
    roi_xmin = w
    roi_xmax = 0
    roi_ymin = h
    roi_ymax = 0
    for i in range(len(contours)-1,-1,-1):
        rect = cv2.boundingRect(contours[i])
        if max(rect[2],rect[3]) >= thresh:
            roi_xmin = min(roi_xmin,rect[0])
            roi_xmax = max(roi_xmax,rect[0]+rect[2])
            roi_ymin = min(roi_ymin,rect[1])
            roi_ymax = max(roi_ymax,rect[1]+rect[3])
    roi = img[roi_ymin:roi_ymax,roi_xmin:roi_xmax]
    return roi

def merge_bgfg(bg_folder,fg_folder,ratio_size=None,ratio_alpha=None):
    '''
    合成烟雾图片
    :param bg:
    :param fg:
    :param ratio:
    :return:
    '''
    #1.处理背景图片，如果size大于2000，则2倍下采样
    if isinstance(bg_folder,str):
        bg_img = cv2.imread(bg_folder, -1)
        h_bg, w_bg, c_bg = bg_img.shape
        assert c_bg == 4 or c_bg == 3, f'{bg_folder} channels={c_bg}!'
        while h_bg >= 2000 or w_bg >= 2000:
            bg_img = cv2.resize(bg_img, (w_bg // 2, h_bg // 2))
            h_bg, w_bg, c_bg = bg_img.shape
        bg_img = cv2.resize(bg_img,(1920,1080))
    else:
        bg_img = bg_folder

    #2.处理烟雾图片
    # fg_img = cv2.imread(fg_folder, -1)
    fg_img = cv_imread(fg_folder,-1)
    # fg_img = filter_smoke_data(fg_img)
    h_fg, w_fg, c_fg = fg_img.shape
    if h_fg==0 or w_fg==0:
        return None,None
    assert c_fg == 4, f'{fg_folder} channels={c_fg}!'
    if ratio_size is None:
        new_h,new_w = h_fg, w_fg
    else:
        new_h,new_w = int(h_fg*ratio_size[0]),int(w_fg*ratio_size[1])
    fg_img = cv2.resize(fg_img,(new_w,new_h))
    h_fg, w_fg, c_fg = fg_img.shape

    #3.合成数据
    #3.1烟雾左上角和右下角坐标
    # tl_x = random.randint(0,w_bg-w_fg)
    # tl_y = random.randint(0,h_bg-h_fg)
    tl_x = 0
    tl_y = 0
    br_x = tl_x + w_fg
    br_y = tl_y + h_fg

    #3.2合成图片
    b, g, r, a = cv2.split(fg_img)

    ratio = 1
    # res_img = np.zeros_like(bg_img)
    res_img = deepcopy(bg_img)

    # inds = a > 0
    # for i in range(3):
    #     res_img[...,i][tl_y:br_y, tl_x:br_x][inds] = fg_img[...,i][inds]
    #
    fg_img[...,:3] = 255 - fg_img[...,:3]
    for i in range(3):
        res_img[...,i][tl_y:br_y, tl_x:br_x] = bg_img[...,i][tl_y:br_y, tl_x:br_x] * (255.0 - ratio * a) / 255
        res_img[...,i][tl_y:br_y, tl_x:br_x] += np.array(fg_img[...,i] * (ratio * a / 255), dtype=np.uint8)

    #3.3烟雾轮廓
    contours,_ = cv2.findContours(a.copy(),cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
    for contour in contours:
        contour[...,0] += tl_x
        contour[...,1] += tl_y

    return res_img,contours

def load_smokedata(smoke_folder='/dataset/smoke_data/2.24'):
    '''
    读取烟雾图片数据
    :param smoke_folder: 烟雾数据根目录
    :return: 返回所有烟雾图片路径对应list
    '''
    dirs = [os.path.join(smoke_folder, dir) for dir in os.listdir(smoke_folder) if not dir.startswith('2.2')]
    list_smoke = []
    for dir in dirs:
        list_bgs = [os.path.join(dir,f) for f in os.listdir(dir)]
        list_smoke += list_bgs
    print(f"smoke data num:{len(list_smoke)}")
    return list_smoke

def load_waterdata():
    '''
    读取烟雾图片数据
    :param smoke_folder: 烟雾数据根目录
    :return: 返回所有烟雾图片路径对应list
    '''
    target_folders = ['正常','明度','模糊']
    target_path = [[],[],[]]
    for i,target_folder in enumerate(target_folders):
        list_files = [os.path.join(target_folder,file) for file in os.listdir(target_folder)]
        target_path[i] = list_files

    return target_path

def load_backgrounds(bg_folder = '/dataset/smoke_data/all_cams'):
    '''
    读取摄像头背景数据
    :param bg_folder:
    :return:
    '''

    list_files = [os.path.join(bg_folder,f) for f in os.listdir(bg_folder)]
    return list_files

if __name__ == '__main__':
    # img_save_folder = '/dataset/smoke_data/synthetic_0318/images'
    # label_save_folder = '/dataset/smoke_data/synthetic_0318/labels'
    # os.makedirs(img_save_folder, exist_ok=True)
    # os.makedirs(label_save_folder, exist_ok=True)
    #加载烟雾数据对应文件路径
    # list_smoke = load_smokedata('G:\\2.24\\2.24/')
    # list_water = load_waterdata()
    list_water = [['2.tif'],['1.tif'],['3.tif']]
    #加载背景数据对应路径
    list_files = load_backgrounds('G:\\1\\3\\4\images\images/')

    save_id = 0
    # bg_folder = random.sample(list_files, 1)[0]
    bg_folder = 'bg2.png'

    length = len(list_water[0])
    for list_data in list_water:
        assert len(list_data)==length

    for i in range(length):
        ratio_size = None
        ratio_alpha = None
        for j in range(len(list_water)):
            smoke_folder = list_water[j][i]
            res_img, contours = merge_bgfg(bg_folder, smoke_folder, ratio_size=ratio_size, ratio_alpha=ratio_alpha)
            if res_img is None:
                continue
            cv2.imshow('img_{}'.format(j), res_img)
            fg_img = cv2.imread(smoke_folder, -1)
            cv2.imshow('smoke', fg_img)
            cv2.waitKey()
            bg_folder = deepcopy(res_img)



